'''
@author: zhangkai
@license: (C) Copyright 2017-2023
@contact: jeffcobile@gmail.com
@Software : PyCharm
@file: deeplabv3_plus.py
@time: 2020-05-29 12:07:15
@desc: 
'''
from ELib.model.model_zoo import MODEL_ZOO
from ELib.basic_model import _ConvBNReLU, _FCNHead
from ELib.model.segbase import SegBaseModel
import torch
from collections import OrderedDict


@MODEL_ZOO.register(name='DeepLabV3_Plus')
class DeepLabV3Plus(SegBaseModel):
    r"""DeepLabV3Plus
    Reference:
        Chen, Liang-Chieh, et al. "Encoder-Decoder with Atrous Separable Convolution for Semantic
        Image Segmentation."
    """
    def __init__(self, cfg):
        super(DeepLabV3Plus, self).__init__(cfg)
        if self.backbone.startswith('mobilenet'):
            c1_channels = 24
            c4_channels = 320
        else:
            c1_channels = 256
            c4_channels = 2048
        self.head = _DeepLabHead(cfg, self.nclass, c1_channels=c1_channels, c4_channels=c4_channels)
        if self.aux:
            self.auxlayer = _FCNHead(728, self.nclass)
        self.__setattr__('decoder', ['head', 'auxlayer'] if self.aux else ['head'])

    def forward(self, x):
        size = x.size()[2:]
        c1, _, c3, c4 = self.encoder(x)

        outputs = list()
        x = self.head(c4, c1)
        x = torch.nn.functional.interpolate(x, size, mode='bilinear', align_corners=True)

        outputs.append(x)
        if self.aux:
            auxout = self.auxlayer(c3)
            auxout = torch.nn.functional.interpolate(auxout, size, mode='bilinear', align_corners=True)
            outputs.append(auxout)
        return tuple(outputs)


class _DeepLabHead(torch.nn.Module):
    def __init__(self, cfg, nclass, c1_channels=256, c4_channels=2048, norm_layer=torch.nn.BatchNorm2d):
        super(_DeepLabHead, self).__init__()
        self.use_aspp = cfg.MODEL.DEEPLABV3_PLUS.USE_ASPP
        self.use_decoder = cfg.MODEL.DEEPLABV3_PLUS.ENABLE_DECODER
        last_channels = c4_channels
        if self.use_aspp:
            self.aspp = _ASPP(cfg, c4_channels, 256)
            last_channels = 256
        if self.use_decoder:
            self.c1_block = _ConvBNReLU(c1_channels, 48, 1, norm_layer=norm_layer)
            last_channels += 48
        self.block = torch.nn.Sequential(
            SeparableConv2d(last_channels, 256, 3, norm_layer=norm_layer, relu_first=False),
            SeparableConv2d(256, 256, 3, norm_layer=norm_layer, relu_first=False),
            torch.nn.Conv2d(256, nclass, 1))

    def forward(self, x, c1):
        size = c1.size()[2:]
        if self.use_aspp:
            x = self.aspp(x)
        if self.use_decoder:
            x = torch.nn.functional.interpolate(x, size, mode='bilinear', align_corners=True)
            c1 = self.c1_block(c1)
            return self.block(torch.cat([x, c1], dim=1))

        return self.block(x)


class _ASPP(torch.nn.Module):
    def __init__(self, cfg, in_channels=2048, out_channels=256):
        super().__init__()
        output_stride = cfg.MODEL.OUTPUT_STRIDE
        if output_stride == 16:
            dilations = [6, 12, 18]
        elif output_stride == 8:
            dilations = [12, 24, 36]
        elif output_stride == 32:
            dilations = [6, 12, 18]
        else:
            raise NotImplementedError

        self.aspp0 = torch.nn.Sequential(OrderedDict([('conv', torch.nn.Conv2d(in_channels, out_channels, 1, bias=False)),
                                                ('bn', torch.nn.BatchNorm2d(out_channels)),
                                                ('relu', torch.nn.ReLU(inplace=True))]))
        self.aspp1 = SeparableConv2d(in_channels, out_channels, dilation=dilations[0], relu_first=False)
        self.aspp2 = SeparableConv2d(in_channels, out_channels, dilation=dilations[1], relu_first=False)
        self.aspp3 = SeparableConv2d(in_channels, out_channels, dilation=dilations[2], relu_first=False)

        self.image_pooling = torch.nn.Sequential(OrderedDict([('gap', torch.nn.AdaptiveAvgPool2d((1, 1))),
                                                        ('conv', torch.nn.Conv2d(in_channels, out_channels, 1, bias=False)),
                                                        ('bn', torch.nn.BatchNorm2d(out_channels)),
                                                        ('relu', torch.nn.ReLU(inplace=True))]))

        self.conv    = torch.nn.Conv2d(out_channels*5, out_channels, 1, bias=False)
        self.bn      = torch.nn.BatchNorm2d(out_channels)
        self.relu    = torch.nn.ReLU(inplace=True)
        self.dropout = torch.nn.Dropout2d(p=0.1)

    def forward(self, x):
        pool = self.image_pooling(x)
        pool = torch.nn.functional.interpolate(pool, size=x.shape[2:], mode='bilinear', align_corners=True)

        x0 = self.aspp0(x)
        x1 = self.aspp1(x)
        x2 = self.aspp2(x)
        x3 = self.aspp3(x)
        x = torch.cat((pool, x0, x1, x2, x3), dim=1)

        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        x = self.dropout(x)

        return x


class SeparableConv2d(torch.nn.Module):
    def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, relu_first=True,
                 bias=False, norm_layer=torch.nn.BatchNorm2d):
        super().__init__()
        depthwise = torch.nn.Conv2d(inplanes, inplanes, kernel_size,
                              stride=stride, padding=dilation,
                              dilation=dilation, groups=inplanes, bias=bias)
        bn_depth = norm_layer(inplanes)
        pointwise = torch.nn.Conv2d(inplanes, planes, 1, bias=bias)
        bn_point = norm_layer(planes)

        if relu_first:
            self.block = torch.nn.Sequential(OrderedDict([('relu', torch.nn.ReLU()),
                                                    ('depthwise', depthwise),
                                                    ('bn_depth', bn_depth),
                                                    ('pointwise', pointwise),
                                                    ('bn_point', bn_point)
                                                    ]))
        else:
            self.block = torch.nn.Sequential(OrderedDict([('depthwise', depthwise),
                                                    ('bn_depth', bn_depth),
                                                    ('relu1', torch.nn.ReLU(inplace=True)),
                                                    ('pointwise', pointwise),
                                                    ('bn_point', bn_point),
                                                    ('relu2', torch.nn.ReLU(inplace=True))
                                                    ]))

    def forward(self, x):
        return self.block(x)
